# 使用torchvision库中的MNIST数据集
import cv2
import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.utils.data as data_utils
import CNN

test_loader = datasets.MNIST(root='mnist', train=False, download=True, transform=transforms.ToTensor())

cnn = torch.load('mnist_cnn.pth')
cnn = cnn.cuda()
loss_test = 0
rightValue = 0
loss_func = torch.nn.CrossEntropyLoss()
for i, (images, labels) in enumerate(test_loader):
    labels = labels.cuda()
    # 前向传播
    outputs = cnn(images.cuda())
    _, predicted = torch.max(outputs.data, 1)
    loss_test += loss_func(outputs, labels).item()
    rightValue += (predicted == labels).sum().item()
    images = images.cpu().numpy()
    labels = labels.cpu().numpy()
    predicted = predicted.cpu().numpy()
    for j in range(images.shape[0]):
        im_data  = images[j]
        im_data = im_data.transpose(1, 2, 0)
        im_label = labels[j]
        im_predicted= predicted[j]
        print("图片id: {}, 预测值: {}, 真实值: {}".format(i * 100 + j, predicted[j], labels[j]))
        cv2.imshow("image", im_data.reshape(28, 28))
        cv2.waitKey(0)

print("Test Loss: {:.6f}, Accuracy: {}".format(loss_test / len(test_loader), rightValue / len(test_loader)))